from .board import GameState, Move, IsPointAnEye
from .gotypes import Point, Player
import numpy as np


class Encoder:
    def Encode(self, game_state:GameState):
        raise NotImplementedError
    
    def EncodePoint(self, point: Point):
        ''' Encode point to index '''
        raise NotImplementedError
    
    def DecodePoint(self, index):
        ''' Decode point from index'''
        raise NotImplementedError
    
    def NumberPoints(self):
        raise NotImplementedError
    
    def Shape(self):
        raise NotImplementedError
    
class OnePlaneEncoder(Encoder):
    def __init__(self, board_size) -> None:
        super().__init__()
        if isinstance(board_size, int):
            board_size = board_size,board_size
        self.board_width, self.board_height = board_size
        self.num_planes = 1
        self.name = 'oneplane'

    def Encode(self, game_state: GameState):
        board_matrix = np.zeros(self.Shape())
        next_player = game_state.next_player
        for r in range(self.board_height):
            for c in range(self.board_width):
                p = Point(r+1, c+1)
                gostring = game_state.board.GetGoString(p)
                if gostring is None:
                    continue
                if gostring.next_player == next_player:
                    board_matrix[0,r,c] = 1
                else:
                    board_matrix[0,r,c] = -1
        return board_matrix
    
    def EncodePoint(self, point: Point):
        return self.board_width * (point.row - 1) + (point.col-1)
    
    def DecodePoint(self, index):
        row = index // self.board_width 
        col = index % self.board_width
        return Point(row+1, col+1)
    
    def NumberPoints(self):
        return self.board_width * self.board_height
    
    def Shape(self):
        return self.num_planes, self.board_height, self.board_width
    
class SimpleEncoder(Encoder):
    def __init__(self, board_size) -> None:
        super().__init__()
        if isinstance(board_size, int):
            board_size = board_size,board_size
        self.board_width, self.board_height = board_size
        self.num_planes = 11
        self.name = 'simple'

    def Encode(self, game_state: GameState):
        board_tensor = np.zeros(self.Shape())
        if game_state.next_player == Player.black:
            board_tensor[8] = 1
        else:
            board_tensor[9] = 1
        for r in range(self.board_height):
            for c in range(self.board_width):
                p = Point(r+1, c+1)
                go_string = game_state.board.GetGoString(p)
                if go_string is None:
                    if game_state.DoesMoveViolateKo(game_state.next_player, Move.Play(p)):
                        board_tensor[10][r][c] = 1
                else:
                    liberty_plane = min(4, go_string.NumberLiberties) - 1
                    if go_string.next_player == Player.white:
                        liberty_plane += 4
                    board_tensor[liberty_plane][r][c] = 1
        return board_tensor
    
    def EncodePoint(self, point: Point):
        return self.board_width * (point.row - 1) + (point.col - 1)
    
    def DecodePoint(self, index):
        row = index // self.board_width
        col = index % self.board_width
        return Point(row+1, col+1)
    
    def NumberPoints(self):
        return self.board_width * self.board_height
    
    def Shape(self):
        return self.num_planes, self.board_height, self.board_width

    
class SevenPlaneEncoder(Encoder):
    def __init__(self, board_size) -> None:
        super().__init__()
        if isinstance(board_size, int):
            board_size = board_size,board_size
        self.board_width, self.board_height = board_size
        self.num_planes = 7
        self.name = 'sevenplane'

    def Encode(self, game_state: GameState):
        board_tensor = np.zeros(self.Shape())
        base_plane = {game_state.next_player: 0, game_state.next_player.Other: 3}
        for r in range(self.board_height):
            for c in range(self.board_width):
                p = Point(r+1, c+1)
                go_string = game_state.board.GetGoString(p)
                if go_string is None:
                    if game_state.DoesMoveViolateKo(game_state.next_player, Move.Play(p)):
                        board_tensor[6][r][c] = 1
                else:
                    liberty_plane = min(3, go_string.NumberLiberties) - 1
                    liberty_plane += base_plane[go_string.next_player]
                    board_tensor[liberty_plane][r][c] = 1
        return board_tensor
    
    def EncodePoint(self, point: Point):
        return self.board_width * (point.row - 1) + (point.col - 1)
    
    def DecodePoint(self, index):
        row = index // self.board_width
        col = index % self.board_width
        return Point(row+1, col+1)
    
    def NumberPoints(self):
        return self.board_width * self.board_height
    
    def Shape(self):
        return self.num_planes, self.board_height, self.board_width

    
class AlphaGoEncoder(Encoder):
    FEATURE_OFFSETS = {
        "stone_color": 0,
        "ones": 3,
        "zeros": 4,
        "sensibleness": 5,
        "turns_since": 6,
        "liberties": 14,
        "liberties_after": 22,
        "capture_size": 30,
        "self_atari_size": 38,
        "ladder_capture": 46,
        "ladder_escape": 47,
        "current_player_color": 48
    }

    def __init__(self, board_size=(19, 19), use_player_plane=False) -> None:
        super().__init__()
        if isinstance(board_size ,int):
            board_size = board_size, board_size
        self.board_width, self.board_height = board_size
        self.use_player_plane = use_player_plane
        self.num_planes = 48 + use_player_plane
        self.name = 'alphago'

    def Offset(self, feature):
        return self.FEATURE_OFFSETS[feature]
    
    def Encode(self, game_state: GameState):
        board_tensor = np.zeros((self.num_planes, self.board_height, self.board_width))
        for r in range(self.board_height):
            for c in range(self.board_width):
                point = Point(r+1, c+1)
                go_string = game_state.board.GetGoString(point)
                if go_string and go_string.next_player == game_state.next_player:
                    board_tensor[self.Offset('stone_color')][r][c] = 1
                elif go_string and go_string.next_player == game_state.next_player.Other:
                    board_tensor[self.Offset('stone_color')+1][r][c] = 1
                else:
                    board_tensor[self.Offset('stone_color')+2][r][c] = 1

                board_tensor[self.Offset('ones')] = self.Ones()
                board_tensor[self.Offset('zeros')] = self.Zeros()

                if not IsPointAnEye(game_state.board, point, game_state.next_player):
                    board_tensor[self.Offset('sensibleness')][r][c] = 1
                ages = min(game_state.board.move_ages.GetMoveAge(r,c), 8)
                if ages > 0:
                    print(ages)
                    board_tensor[self.Offset('turns_since')][r][c] = 1
                if game_state.board.GetGoString(point):
                    liberities = min(game_state.board.GetGoString(point).NumberLiberties, 8)
                    board_tensor[self.Offset('liberties') + liberities][r][c] = 1
                
                move = Move(point)
                if game_state.IsValidMove(move):
                    new_state = game_state.ApplyMove(move)
                    liberities = min(new_state.board.GetGoString(point).NumberLiberties, 8)
                    board_tensor[self.Offset('liberties_after') + liberities][r][c] = 1
                    adjacent_strings = [game_state.board.GetGoString(nb) for nb in point.Neighbors()]
                    capture_count = 0
                    for s in adjacent_strings:
                        other_player = game_state.next_player.Other
                        if s and s.NumberLiberties == 1 and s.next_player == other_player:
                            capture_count += len(s.stones)
                    capture_count = min(capture_count, 8)
                    board_tensor[self.Offset('capture_size') + capture_count][r][c] = 1
                if go_string and go_string.NumberLiberties == 1:
                    go_string = game_state.board.GetGoString(point)
                    if go_string:
                        num_atari_stones = min(len(go_string.stones), 8)
                        board_tensor[self.Offset('self_atari_size') + num_atari_stones][r][c] = 1
                if IsLadderCapture(game_state, point):pass


    def Ones(self):
        return np.ones((1, self.board_height, self.board_width))
    
    def Zeros(self):
        return np.zeros((1, self.board_height, self.board_width))
    
    def EncodePoint(self, point: Point):
        return self.board_width * (point.row -1) + (point.col -1)
    
    def DecodePoint(self, index):
        row = index // self.board_width
        col = index % self.board_width
        return Point(row+1, col+1)
    
    def NumberPoints(self):
        return self.board_width *self.board_height
    
    def Shape(self):
        return self.num_planes, self.board_height, self.board_width
    
def IsLadderCapture(game_state:GameState, candidate:Point, recursion_depth=50):
    return IsLadder(True, game_state, candidate, None, recursion_depth)

def IsLadderEscape(game_state, candidate:Point, recursion_depth=50):
    return IsLadder(False, game_state, candidate, None, recursion_depth)

def IsLadder(try_capture, game_state:GameState, candidate:Point, ladder_stones=None, recursion_depth=50):
    if not game_state.IsValidMove(Move(candidate)) or not recursion_depth:
        return False
    next_player = game_state.next_player
    caoture_player = next_player if try_capture else next_player.Other
    escape_player = caoture_player.Other
    if ladder_stones is None:
        ladder_stones = GuessLadderStones(game_state, candidate, escape_player)
    for ladder_stone in ladder_stones:
        current_state = game_state.ApplyMove(Move(candidate))
        if try_capture:
            candidates = DetermineEscapeCandidates(game_state, ladder_stone, caoture_player)
            attempted_escapes = [IsLadder(False, current_state, escape_candidate, ladder_stone, 
                                          recursion_depth -1) 
                                for escape_candidate in candidates]
            if not any(attempted_escapes):
                return True 
        else:
            if CountLiberities(current_state, ladder_stone) >= 3:
                return True 
            if CountLiberities(current_state, ladder_stone) == 1:
                continue
            candidates = Liberities(current_state, ladder_stone)
            attempted_captures = [
                IsLadder(True, current_state, capture_candidate, ladder_stone, 
                         recursion_depth -1)
                for capture_candidate in candidates
            ]
            if any(attempted_captures):
                continue
            return True 
    return False

def DetermineEscapeCandidates(game_state:GameState, point:Point, player:Player):
    escape_candidates = point.Neighbors()
    for other_ladder_stone in game_state.board.GetGoString(point).stones:
        for neighbor in other_ladder_stone.Neighbors():
            right_player = game_state.board.GetPlayer(neighbor) == player
            one_liberity = CountLiberities(game_state, neighbor) == 1
            if right_player and one_liberity:
                escape_candidates.append(Liberities(game_state, neighbor))
    return escape_candidates

def Liberities(game_state:GameState, point:Point):
    return list(game_state.board.GetGoString(point).liberties)

def GuessLadderStones(game_state:GameState, point:Point, escape_player:Player):
    adjacent_strings = [game_state.board.GetGoString(nb) for nb in point.Neighbors() 
                        if game_state.board.GetGoString(nb)]
    if adjacent_strings:
        gostring = adjacent_strings[0]
        neighbors = []
        for s in adjacent_strings:
            stones = s.stones
            for stone in stones:
                neighbors.append(stone)
        return [nb for nb in neighbors if IsCandidate(game_state, nb, escape_player)]
    else:
        return []
    
def IsCandidate(game_state:GameState, point:Point, player:Player):
    return game_state.next_player == player and CountLiberities(game_state, point) == 2

def CountLiberities(game_state:GameState, point:Point):
    if game_state.board.GetGoString(point):
        return game_state.board.GetGoString(point).NumberLiberties
    else:
        return 0